# coding=utf-8
import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))


from chains.C2_recall_material import recall_material_list, recall_material_distances_db
from chains.C3_material_problem import Material2Problem
from chains.C4_get_prompt_v2 import get_prompt, get_material_category
from chains.C5_AI_answer import AI_Answer
from chains.C20_get_prompt_third import get_prompt_third_db
from chains.C13_bubble_score import Global_Score
from chains.C15_bubble_title import Bubble
from prompts.C24_prompt import prompt_24
from config.config import config
from collections import OrderedDict
from utils.functions import Get_Dict
from datetime import datetime
# from utils.db import LawDB
from utils.db_pool import LawDB
from pprint import pprint
from openai import OpenAI, AsyncOpenAI
import aiomysql
import asyncio
import time
import math
import json
from utils.utils import remove_json_prefix_suffix


"""
该版本：
"""


class Overall_Bubble_V2:

    def __init__(self):

        self.client = AsyncOpenAI(
            # api_key = "575037f0-0614-4f49-a142-fdb830abfba3",
            api_key="5db64dc0-73d7-4139-b755-5bd5483327b4",  # lz
            base_url="https://ark.cn-beijing.volces.com/api/v3",
            max_retries=3,
            timeout=30,
        )

    async def C2_distance(self, db: LawDB, content_text_dict: dict, sector: str, industry: str, region: str):
        tasks = set()
        for index, content_text in content_text_dict.items():
            distances_bubble_task = asyncio.create_task(
                self.single_slice_distances(db, content_text, index, sector, industry, region))
            tasks.add(distances_bubble_task)
            # 添加回调函数，在任务完成后从集合中移除
            distances_bubble_task.add_done_callback(tasks.discard)
            # 等待所有任务完成
        results = await asyncio.gather(*tasks)
        distances_ordered_dict = {}
        for result in results:
            distances_ordered_dict.update(result)
        return distances_ordered_dict

    # 切片存入数据库
    async def single_slice_distances(self, db: LawDB, content_text: str, ordered_index: int, sector: str, industry: str, region: str):

        time_0 = time.time()
        Title = Bubble()
        problem_type_list = await Title.title_bubble_v3(content_text, sector)
        time_1 = time.time()

        print('调用 C15 获取问题分类标签 时间：', time_1 - time_0)

        try:
            await db.init_pool()
            distances_list = await recall_material_distances_db(db.pool, content_text, 3, sector, industry, region, problem_type_list)

            if len(distances_list) == 0:
                print('C2 列表长度=0，返回空字典:', distances_list)
                return {}
            elif isinstance(distances_list, list):
                pass
            else:
                print('type(distances_list)', type(distances_list))
                print('C2 输出是字符串不是列表 distances_list:', distances_list)
                return {}
        except Exception as e:
            print('C2 输出是字符串不是列表返回空字典:', e)
            return {}

        time_2 = time.time()
        print('调用 C2 获得材料详情列表 时间：', time_2 - time_1)

        material_id_list = []
        material_distances_list = []
        material_content_list = []

        try:
            for material_dict in distances_list:
                print('material_dict:', material_dict)
                material_id_list.append(list(material_dict.keys())[0])
                material_distances_list.append(list(material_dict.values())[0]['材料距离'])
                material_content_list.append(list(material_dict.values())[0]['材料内容'])
            min_distance = min(material_distances_list).item()
            print('最小距离：', min_distance)
        except Exception as e:
            raise ValueError(f"distances_list_type:{type(distances_list)};distances_list: {distances_list};content_text: {content_text};problem_type_list: {problem_type_list};exception: {e}")

        bubble_dict = {}
        bubble_dict['contentText'] = content_text
        bubble_dict['bubbleTitle'] = problem_type_list
        bubble_dict['materials'] = material_content_list
        bubble_dict['material_id_list'] = material_id_list


        return {ordered_index: {'min_distance': min_distance, 'bubble_dict': bubble_dict}}

    # 处理 content_text_dict（一批次的招股书原文字典） 并输出 review_dict_asc（需要产生气泡的字典）
    async def R1_score(self, content_text_dict: dict):

        s_time = time.time()
        print("----- streaming request -----")
        query_message = prompt_24(content_text_dict)
        # print('query_message')
        # print(query_message)

        stream = await self.client.chat.completions.create(
        # stream = self.client.chat.completions.create(
            model="ep-20250210135803-8r8cd",  # lz 满血
            # model="ep-20250210135919-kncgs",  # lz 32b
            messages=[
                {"role": "user", "content": query_message},
            ],
            stream=True,
            timeout=300,
        )
        reasoning_content = ""
        content = ""
        async for chunk in stream:
        # for chunk in stream:
            if hasattr(chunk.choices[0].delta, 'reasoning_content'):
                # print(chunk.choices[0].delta.reasoning_content, end='')
                reasoning_content += chunk.choices[0].delta.reasoning_content
            else:
                # print(chunk.choices[0].delta.content, end='')
                content += chunk.choices[0].delta.content
        print('content', content)
        score_dict = eval(Get_Dict(content))
        # 如果需要保持有序，可以使用 OrderedDict
        ordered_dict = OrderedDict(sorted(score_dict.items(), key=lambda item: item[1], reverse=True))

        # 将键从字符串转换为整数
        R1_dict = {int(key): value for key, value in ordered_dict.items()}
        return R1_dict


    async def single_bubble(self, db: LawDB, review_index: int, review_score: int, user_id: int,
                            fileId: str, percentage: float, content_slice_dict: dict):
        content_text = content_slice_dict[str(review_index)]
        bubble_dict = {}
        try:
            await db.init_pool()
            bubble = await db.get_history_overall_bubble(fileId, review_index)
        except Exception as e:
            print('获取 bubbles 出错， 删除切片气泡')
            await db.init_pool()
            await db.delete_history_overall_bubble(fileId, review_index)
        # 因为这里都插入了，所以直接为空
        print('bubble type', type(bubble))
        bubble_dict = bubble
        # 找到相似的材料
        material_id_list = None
        if bubble_dict:
            material_id_list = bubble_dict.get('material_id_list', None)
        Material = Material2Problem()

        problem_id_list = None
        if material_id_list:
            problem_id_list = await Material.material_problem_db(db.pool, material_id_list)

        # 根据问询规则召回对应的prompt
        category_path = os.path.join(config.DATA_DIR, 'rules_sample.json')
        rules_path = os.path.join(config.DATA_DIR, 'rules_query_merged.json')

        time_4 = time.time()
        AI = AI_Answer()
        aiText = ''
        aiText_think = ''

        category_list = await get_material_category(AI, content_text, category_path)
        aiText_prompt = await get_prompt(AI, content_text, category_list, rules_path)

        async for Text in AI.get_ai_answer_R1(aiText_prompt, output_json=True):
            Text_Type = json.loads(Text)
            if Text_Type['type'] == 'thinking':
                aiText_think += Text_Type['content']
            else:
                aiText += Text_Type['content']
        print('review_index: %s, aiText: %s\n' % (review_index, aiText))
        aiText = remove_json_prefix_suffix(aiText)
        time_5 = time.time()
        print('调用 C4 获取 发行人 回答 时间：', time_5 - time_4)
        if not bubble_dict:
            bubble_dict = {}
        bubble_dict_aiText_list = []
        for index, inquiry_item in enumerate(aiText['inquiry']):
            temp_inquiry_item = f'({index+1}) {inquiry_item};'
            bubble_dict_aiText_list.append(temp_inquiry_item)
        bubble_dict['aiText'] = '\n'.join(bubble_dict_aiText_list)
        bubble_dict['aiText_think'] = aiText_think
        bubble_dict['rule_ref'] = aiText['rules_ref']
        bubble_dict['useFlag'] = False
        bubble_dict['xialaCut'] = False
        bubble_dict['aiBtnFlag'] = False
        bubble_dict['cutBtnFlag'] = 2
        bubble_dict['sliceIndex'] = review_index
        bubble_dict['bubbleScore'] = review_score
        bubble_dict['contentText'] = content_text
        bubble_dict['domElementCur'] = None
        bubble_dict['problem_id_list'] = problem_id_list


        await db.init_pool()
        aiTextThird_prompt = await get_prompt_third_db(db.pool, content_text, problem_id_list, aiText)
        aiTextThird = ''
        aiTextThird_think = ''
        async for Text in AI.get_ai_answer_R1(aiTextThird_prompt):
            Text_Type = json.loads(Text)
            if Text_Type['type'] == 'thinking':
                aiTextThird_think += Text_Type['content']
            else:
                aiTextThird += Text_Type['content']
        time_6 = time.time()
        print('调用 C5 获取 第三方 回答 时间：', time_6 - time_5)

        bubble_dict['aiTextThird'] = aiTextThird
        bubble_dict['aiTextThird_think'] = aiTextThird_think

        # db = LawDB()
        # await db.init_conn()
        # 初始化连接池
        await db.init_pool()
        await db.update_history_overall(fileId=fileId, bubble=bubble_dict, status_code=0,
                                        slice_index=int(review_index),
                                        bubble_num=float(percentage))
        print(bubble_dict)
        print('单气泡产生完成')


    # 处理 review_dict_asc 并输出AI回复
    async def bubbles_generated(self, R1_dict: dict, C1_dict: dict, db: LawDB, user_id: int, file_name: str, fileId: str, percentage: float, content_slice_dict: dict):
        """
        :param R1_dict:
        :param C1_dict: 原始版本利用作为相似材料
        :param db:
        :param user_id:
        :param file_name:
        :param fileId:
        :param percentage:
        :param content_slice_dict:
        :return:
        """
        review_dict = {}
        for R1_index, R1_score in R1_dict.items():
            # 问询可能性小于7分的不处理
            if R1_score <= 7:
                continue
            review_dict[int(R1_index)] = R1_score

            # v1根据相似材料，进行问询。v2基于规则问询，但是如果存在相似材料，需要将相似材料插入数据库, 不存在相似材料，直接插入空值
            C1_distance = C1_dict.get(str(R1_index))
            bubble = C1_distance['bubble_dict'] if C1_distance else None
            distance = C1_distance['min_distance'] if C1_distance else None
            await db.init_pool()
            await db.insert_history_overall(user_id=user_id, file_name=file_name, bubble=bubble, fileId=fileId,
                                             status_code=0, slice_index=int(R1_index), score=R1_score, distance = distance)

        # 按键升序排序
        review_dict_asc = dict(sorted(review_dict.items(), key=lambda item: item[0]))
        print('review_dict_asc:', review_dict_asc)

        print('该批次有', len(review_dict_asc), '个气泡需要产生')
        tasks = set()
        for num, (review_index, review_score) in enumerate(review_dict_asc.items()):
            single_bubble_task = asyncio.create_task(self.single_bubble(db, review_index, review_score, user_id, fileId, percentage, content_slice_dict))
            tasks.add(single_bubble_task)
            # 添加回调函数，在任务完成后从集合中移除
            single_bubble_task.add_done_callback(tasks.discard)
            # 等待所有任务完成
        await asyncio.gather(*tasks)
        print('单批次气泡产生完成')

    # 主函数
    async def overall_bubble(self, index_paragraph_sentence_dict: dict, sector: str, industry: str, region: str, user_id: int, file_name: str, fileId: str):
        tasks = set()
        json_len = len(index_paragraph_sentence_dict)
        print("切片JSON长度", json_len)

        db = LawDB()
        try:
            # 初始化连接池
            await db.init_pool()

            key_value_list = list(index_paragraph_sentence_dict.items())  # 将字典转换为键值对的列表
            start_time = time.time()
            for i in range(0, len(key_value_list), 16):
                print('json_len', json_len)
                batch = key_value_list[i: i + 16]
                type_dict = dict(batch)
                print('分子', (i + len(type_dict)), '分母', json_len)
                print('进度', round((i + len(type_dict) - 1) / json_len, 2))
                percentage = round((i + len(type_dict) - 1) / json_len, 2)

                content_slice_dict = OrderedDict()
                content_text_dict = OrderedDict()
                for k, v in type_dict.items():
                    content_text_dict[k] = v['文本']
                    content_slice_dict[k] = v['切片']
                print(len(content_text_dict))
                print(list(content_text_dict.keys()))

                # 启动 function1 并获取 review_dict_asc
                # R1_task: 获取文本是否需要问询的分数
                R1_task = asyncio.create_task(self.R1_score(content_text_dict))
                C2_task = asyncio.create_task(self.C2_distance(db, content_text_dict, sector, industry, region))
                # 启动 function2 并传入 review_dict_asc
                ai_reply_task = asyncio.create_task(self.bubbles_generated(await R1_task, await C2_task, db, user_id, file_name, fileId, percentage, content_slice_dict))
                tasks.add(ai_reply_task)
                # 添加回调函数，在任务完成后从集合中移除
                ai_reply_task.add_done_callback(tasks.discard)
            # 等待所有任务完成
            await asyncio.gather(*tasks)
            # db = LawDB()
            # await db.init_conn()
            # 初始化连接池
            await db.init_pool()
            await db.save_overall_progress_bar(fileId=fileId, score=8, distance=0.32)
        finally:
            await db.close_pool()
        # 计算总耗时
        end_time = time.time()
        total_time = end_time - start_time
        print(f"总耗时：{total_time} 秒")



if __name__ == '__main__':
    import asyncio

    Overall = Overall_Bubble_V2()
    test_doc_path = os.path.join(config.PROJECT_DIR, 'data', 'liqi_test_doc.json')
    with open(test_doc_path, 'r', encoding='utf-8') as f:
      index_paragraph_sentence_dict = json.load(f)

    Sector = '科创板'
    Region = ""
    user_id = 1
    file_name = '屹唐股份0806.docx'
    file_id = 'jp_yitanggufen_0807_test'
    Industry = "计算机、通信和其他电子设备制造业"
    # 获取当前时间
    current_time = datetime.now()
    # 打印当前时间
    print(f"总程序开始时间", current_time)

    loop = asyncio.new_event_loop()
    asyncio.set_event_loop(loop)

    loop.run_until_complete(Overall.overall_bubble(index_paragraph_sentence_dict, Sector, Industry, Region, user_id, file_name, file_id))

    # 获取当前时间
    current_time = datetime.now()

    # 打印当前时间
    print(f"总程序结束时间", current_time)

    # 运行主函数
    # asyncio.run(main())

